/*
 * This program scales a CT image and converts it to an attenuation MUMAP.
 * 
 * The imaging pipeline is as follows:
 * 1. Apply a gaussian filter which is equal to the crystal size
 * 2. Resample the image to match the dimensions of the PET scanner
 * 3. Apply bilinear scaling
 * 4. Flip or rotate the mu-map if necessary
 * 5. Write an Interfile Header which is understood by STIR
 *
 * This was written by Robbie Barnett rbar9508@uni.sydney.edu.au
 */

#include "petac.h"


int main( int argc, char* argv[] )
{
	
	if( argc < 4 )
    {
		std::cerr << "Usage: " << std::endl;
		std::cerr << argv[0] << " <scanner filename> <input filename> <output filename> [ <x-offset> <y-offset> <starting-slice> ]" 
		<< std::endl;
		return EXIT_FAILURE;
    }
	
	// Scanner definition	
	std::ifstream scannerFile;
	scannerFile.open (argv[1]);
	if (! scannerFile.good()) {
		std::cerr << "Scanner file not found: " << argv[1] << std::endl;
		return EXIT_FAILURE;
	}
	std::string line;
	scannerFile >> line;
	int number_of_bins = atoi(line.c_str());
	scannerFile >> line;
	int number_of_rings = atoi(line.c_str());
	scannerFile >> line;
	float bin_size = atof(line.c_str()); // millimetres
	scannerFile >> line;
	float distance_between_rings = atof(line.c_str()); // millimetres	
	scannerFile >> line;	
	int flipX = atoi(line.c_str()) > 0;
	scannerFile >> line;	
	int flipY = atoi(line.c_str()) > 0;
	scannerFile >> line;	
	int flipZ = atoi(line.c_str()) > 0;

	if (! scannerFile.good()) {
		std::cerr << "Warning: Scanner file appears to be truncated please check that the following scanner defintion is correct:" << std::endl;
	}
	
	scannerFile.close();
	
	std::cout << "Scanner definition:" << std::endl;
	std::cout << "Number of bins: " << number_of_bins << ", Number of rings: " << number_of_rings << std::endl;
	std::cout << "Bin size: " << bin_size << "mm, Distance between rings: " << distance_between_rings << "mm" << std::endl;
	std::cout << "Flip axis: " << flipX << ", " << flipY << ", "  << flipZ << std::endl;

	if ((number_of_bins < 1) or (number_of_rings < 1) or (bin_size < ADD_FILTER_DELTA) or (distance_between_rings < ADD_FILTER_DELTA)) {
		std::cerr << "Scanner definition invalid: " << argv[1] << std::endl;
		return EXIT_FAILURE;
	}
	
	
	try {
		ReaderType::Pointer reader = ReaderType::New();
		reader->SetFileName( argv[2] );
		try
		{
			reader->Update();
		}
		catch (itk::ExceptionObject &ex)
		{
			std::cout << ex << std::endl;
			return EXIT_FAILURE;
		}
		typedef itk::ImageFileWriter< ImageType > WriterType;
		WriterType::Pointer writer = WriterType::New();
		
		std::string outputPrefix(argv[3]);
		writer->SetFileName( outputPrefix +   ".mhd");
		
		// Create the gaussian filters in the X and Y directions
		
		typedef itk::RecursiveGaussianImageFilter< 
		ImageType,
		ImageType > GaussianFilterType;
		
		GaussianFilterType::Pointer smootherX = GaussianFilterType::New();
		GaussianFilterType::Pointer smootherY = GaussianFilterType::New();
		
		smootherX->SetInput( reader->GetOutput() );
		smootherY->SetInput( smootherX->GetOutput() );
		
		ImageType::ConstPointer inputImage = reader->GetOutput();
		
		
		smootherX->SetSigma( bin_size );
		smootherY->SetSigma( bin_size );
		
		smootherX->SetDirection( 0 );
		smootherY->SetDirection( 1 );
		
		smootherX->SetOrder( GaussianFilterType::ZeroOrder );
		smootherY->SetOrder( GaussianFilterType::ZeroOrder );
		
		// https://mri.radiology.uiowa.edu/mediawiki/index.php/Gaussian_Filter_Example
		// SetNormalizeAcrossScale- Set/Get the flag for normalizing the gaussian over scale space.
		// When this flag is ON the filter will be normalized in such a way
		// that larger sigmas will not result in the image fading away.
		
		smootherX->SetNormalizeAcrossScale( false );
		smootherY->SetNormalizeAcrossScale( false );
		
		
		reader -> Update();
		
		const ImageType::SpacingType& inputSpacing = reader->GetOutput()->GetSpacing();
		
		// Determine the region which covers the entire CT
		
		ImageType::RegionType inputRegion = reader->GetOutput()->GetLargestPossibleRegion();	
		ImageType::SizeType inputSize = inputRegion.GetSize(); 
		ImageType::PointType inputOrigin = reader->GetOutput()->GetOrigin();
		
		ImageType::PointType outputOrigin;
		outputOrigin[0] = inputOrigin[0];
		outputOrigin[1] = inputOrigin[1];
		outputOrigin[2] = inputOrigin[2];
		
		std::cout << "CT origin is " << outputOrigin[0] << "mm, " << outputOrigin[1] << "mm, " << outputOrigin[2] << "mm " << std::endl;
		std::cout << "CT size is " << inputSize[0] << ", " << inputSize[1] << ", " << inputSize[2] << std::endl;
		std::cout << "CT spacing is " << inputSpacing[0] << "mm, " << inputSpacing[1] << "mm, " << inputSpacing[2] << "mm" << std::endl;
		
		// The initial assumption is that the start of the CT matches the start of the PET.
		// You manually adjust the starting slic corresponding to the appropriate bed position.
		// We also assume that the CT and PET isocentre are the same. A manual shift in the
		// x and y offset can be applied of the x and y offset has been changed
		if( argc >= 7 ) {
			std::cout << "MuMap origin would have been: " << outputOrigin[0] << "mm, " << outputOrigin[1] << "mm, " << outputOrigin[2] << "mm " << std::endl;
			outputOrigin[0] += atof(argv[4]); // Manually shift the xoffset
			outputOrigin[1] += atof(argv[5]); // Manually shift the yoffset
			outputOrigin[2] = atof(argv[6]);
			std::cout << "MuMap origin has been shifted. " << std::endl;
		}
		
		// Create a filter to resample the CT
		
		typedef itk::ResampleImageFilter<
		ImageType, ImageType >  ResampleFilterType;
		ResampleFilterType::Pointer resampler = ResampleFilterType::New();
		
		typedef itk::IdentityTransform< double, Dimension >  TransformType;
		TransformType::Pointer transform = TransformType::New();
		transform->SetIdentity();
		resampler->SetTransform( transform );
		
		typedef itk::LinearInterpolateImageFunction< 
		ImageType, double >  InterpolatorType;
		
		InterpolatorType::Pointer interpolator = InterpolatorType::New();
		
		resampler->SetInterpolator( interpolator );
		
		resampler->SetDefaultPixelValue( -1000 ); // Default pixel value to air ( -1000HU )
		
		typedef ImageType::SizeType::SizeValueType SizeValueType;
		ImageType::SpacingType outputSpacing;
		ImageType::SizeType   outputSize;
		
		// Set the resolution equal to the crystal size
		
		outputSpacing[0] = bin_size;
		outputSpacing[1] = bin_size;
		outputSpacing[2] = distance_between_rings/2;
		
		resampler->SetOutputSpacing( outputSpacing );
		
		// Set the size of the MU map to cover the same volume as the CT
		
		double dx = inputSize[0] * inputSpacing[0] / outputSpacing[0];
		double dy = inputSize[1] * inputSpacing[1] / outputSpacing[1];
		double dz = 2*number_of_rings-1;
		
		int pet_fov = number_of_bins+1;
		
		// Expand the size of the MU map to over the entire field of view
		if (dx*outputSpacing[0] < pet_fov*bin_size) {
			std::cout << "MuMap X size too small: " << dx << std::endl;
			outputOrigin[0] += (dx*outputSpacing[0] - pet_fov*bin_size)/2.0;
			dx = pet_fov*bin_size/outputSpacing[0];
			std::cout << "MuMap X size expanded to: " << dx << std::endl;
		}
		if (dy*outputSpacing[1] < pet_fov*bin_size) {
			std::cout << "MuMap Y size too small: " << dy << std::endl;
			outputOrigin[1] += (dy*outputSpacing[1] - pet_fov*bin_size)/2.0;
			dy = pet_fov*bin_size/outputSpacing[1];
			std::cout << "MuMap Y size expanded to: " << dy << std::endl;
		}
		
		resampler->SetOutputOrigin( outputOrigin );
		
		outputSize[0] = static_cast<SizeValueType>( dx );
		outputSize[1] = static_cast<SizeValueType>( dy );
		outputSize[2] = static_cast<SizeValueType>( dz );
		
		if (outputOrigin[0] > inputOrigin[0] + inputSpacing[0]*inputSize[0]) {
            std::cerr << "Error: MU Map region is off the edge of the CT (X output origin too big)" << std::endl;
			return 2;
        }
		if (outputOrigin[1] > inputOrigin[1] + inputSpacing[1]*inputSize[1]) {
            std::cerr << "Error: MU Map region is off the edge of the CT (Y output origin too big)" << std::endl;
        return 2;
        }
        if (outputOrigin[2] > inputOrigin[2] + inputSpacing[2]*inputSize[2]) {
            std::cerr << "Error: MU Map region is off the edge of the CT (Z output origin too big)" << std::endl;
			return 2;
        }
        if (outputOrigin[0]+outputSpacing[0]*outputSize[0] < inputOrigin[0]) {
            std::cerr << "Error: MU Map region is off the edge of the CT (X output endpoint too small)" << std::endl;
            return 2;
        }
        if (outputOrigin[1]+outputSpacing[1]*outputSize[1] < inputOrigin[1]) {
            std::cerr << "Error: MU Map region is off the edge of the CT (Y output endpoint too small)" << std::endl;
            return 2;
        }
        if (outputOrigin[2]+outputSpacing[2]*outputSize[2] < inputOrigin[2]) {
            std::cerr << "Error: MU Map region is off the edge of the CT (Z output endpoint too small)" << std::endl;
            return 2;
        }
		
		std::cout << "MuMap origin is " << outputOrigin[0] << "mm, " << outputOrigin[1] << "mm, " << outputOrigin[2] << "mm " << std::endl;
		std::cout << "MuMap size is " << outputSize[0] << ", " << outputSize[1] << ", " << outputSize[2] << std::endl;
		std::cout << "MuMap spacing: " << outputSpacing[0] << "mm, " << outputSpacing[1] << "mm, " << outputSpacing[2] << "mm " << std::endl;
		
		resampler->SetSize(outputSize );
		
		resampler->SetInput( smootherY->GetOutput()  );
		
		// Apply bilienar scale to convert hounsfield units to linear attenuation coefficients
		
		// Create two scaled images. The first image scales tissue, the second scales bone.
		ShiftScaleFilterType::Pointer shiftFilter1 = ShiftScaleFilterType::New(); 
		ShiftScaleFilterType::Pointer shiftFilter2 = ShiftScaleFilterType::New(); 
		shiftFilter1->SetInput(resampler->GetOutput());
		shiftFilter2->SetInput(resampler->GetOutput());
		
		// Tissue equivalent scaling
		shiftFilter1->SetShift(shiftScale1[0]);
		shiftFilter1->SetScale(shiftScale1[1]);
		
		// Bone equivalent scaling
		shiftFilter2->SetShift(shiftScale2[0]);
		shiftFilter2->SetScale(shiftScale2[1]);
		
		// Prepare images for adding back together by applying a threshold
		ThresholdFilterType::Pointer thresholdFilter1 = ThresholdFilterType::New(); 
		ThresholdFilterType::Pointer thresholdFilter2 = ThresholdFilterType::New(); 
		ThresholdFilterType::Pointer thresholdFilter2u = ThresholdFilterType::New(); 
		thresholdFilter1->SetInput( shiftFilter1->GetOutput() ); 
		thresholdFilter2->SetInput( shiftFilter2->GetOutput() ); 
		thresholdFilter2u->SetInput( thresholdFilter2->GetOutput() ); 
		
		// Threshold so that we see tissue only, everything else is zero
		thresholdFilter1->SetOutsideValue( 0 ); 
		thresholdFilter1->ThresholdAbove( bilinear_scale_mu[1] ); 
		
		// Threshold so that we see bone only, everything else is zero
		thresholdFilter2->SetOutsideValue( 0 ); 
		thresholdFilter2->ThresholdBelow( bilinear_scale_mu[1] + ADD_FILTER_DELTA); 
		
		// Do not set values greater than the maximum possible mu value
		thresholdFilter2u->SetOutsideValue( bilinear_scale_mu[2] ); 
		thresholdFilter2u->ThresholdAbove( bilinear_scale_mu[2] ); 
		
		// Add images back together
		AddFilterType::Pointer addFilter = AddFilterType::New(); 
		addFilter->SetInput1( thresholdFilter1->GetOutput() ); 
		addFilter->SetInput2( thresholdFilter2u->GetOutput() ); 
		
		// Flip about some axis so that the PET and CT point in the same direction
		// 180 degree rotation may be necessary for wholebody scans where patient is lying with head pointing our of scanner.
		
		FlipType::Pointer flipFilter = FlipType::New();
		FlipType::FlipAxesArrayType axes;
		axes[0] = flipX > 0;
		axes[1] = flipY > 0;
		axes[2] = flipZ > 0;
		
		flipFilter -> SetFlipAxes(axes);
		
		flipFilter -> SetInput(addFilter->GetOutput());	
		writer->SetInput( flipFilter->GetOutput() );
		
		ImageType::PointType writeOrigin = flipFilter->GetOutput()->GetOrigin();
		
		
		std::cout  << "Writing the MuMap to file " << outputPrefix << ".mhd" << std::endl;
		try
		{
			writer->Update();
		}
		catch (itk::ExceptionObject &ex)
		{
			std::cout << ex << std::endl;
			return EXIT_FAILURE;
		}
		
		std::cout  << "Writing the MuMap to interfile (scanner coordinates) " << outputPrefix << ".hv" << std::endl;
		
		// Write interfile header in the PET scanner coordinate system
		std::ofstream   iout;
		iout.open( (outputPrefix + ".hv").c_str() );
		if( iout.fail() )
		{
			std::cerr << "Error writing to Interfile." << std::endl;
			return 2;
		}
		
		// The MU map is in the PET scanner coordinate system where x=0 and y=0 is at the isocentre.
		// z=0 is the position of the first ring
		writeOrigin[0] = -outputSpacing[0]*(outputSize[0]-1)/2.0;
		writeOrigin[1] = -outputSpacing[1]*(outputSize[1]-1)/2.0;
		writeOrigin[2] = 0;
		
		ImageType::RegionType writeRegion = writer->GetInput()->GetLargestPossibleRegion();	
		ImageType::SizeType writeSize = writeRegion.GetSize(); 
		const ImageType::SpacingType& writeSpacing = writer->GetInput()->GetSpacing();
		
		iout << "!INTERFILE := " << std::endl;
		iout << "name of data file := " << outputPrefix << ".raw " << std::endl;
		iout << "!GENERAL DATA :=  " << std::endl;
		iout << "!GENERAL IMAGE DATA := " << std::endl; 
		iout << "!type of data := PET " << std::endl;
		iout << "imagedata byte order := LITTLEENDIAN " << std::endl;
		iout << "!PET STUDY (General) :=  " << std::endl;
		iout << "!PET data type := Image " << std::endl;
		iout << "process status := Reconstructed " << std::endl;
		iout << "!number format := float " << std::endl;
		iout << "!number of bytes per pixel := 4 " << std::endl;
		iout << "number of dimensions := 3 " << std::endl;
		iout << "matrix axis label [1] := x " << std::endl;
		iout << "!matrix size [1] := " << writeSize[0] << std::endl;
		iout << "scaling factor (mm/pixel) [1] := " << writeSpacing[0] << std::endl;
		iout << "matrix axis label [2] := y " << std::endl;
		iout << "!matrix size [2] := " << writeSize[1] << std::endl;
		iout << "scaling factor (mm/pixel) [2] := " << writeSpacing[1] << std::endl;
		iout << "matrix axis label [3] := z " << std::endl;
		iout << "!matrix size [3] := " << writeSize[2] << std::endl;
		iout << "scaling factor (mm/pixel) [3] := " << writeSpacing[2] << std::endl;
		iout << "first pixel offset (mm) [1] := " << writeOrigin[0] << std::endl;
		iout << "first pixel offset (mm) [2] := " << writeOrigin[1] << std::endl;
		iout << "first pixel offset (mm) [3] := " << writeOrigin[2] << std::endl;
		iout << "number of time frames := 1 " << std::endl;
		iout << "image scaling factor[1] := 1 " << std::endl;
		iout << "data offset in bytes[1] := 0 " << std::endl;
		iout << "quantification units := 1 " << std::endl;
	}
	
	
	
	catch (itk::ExceptionObject &ex)
    {
		std::cout << ex << std::endl;
		return EXIT_FAILURE;
    }
	
	return EXIT_SUCCESS;
	
}
